/*
    Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/**  \file
    Declarations of the grid_base and adaptive_grid classes. This file
    includes all the other grid_xxx.h files.  Never include them
    directly. */
// $Id$



#ifndef __grid__
#define __grid__

#include "mcrx-types.h"
#include "mcrx-debug.h"
//

#include "blitz/numinquire.h"
#include <ostream>
#include <vector> 
#include "ray.h"
#include <algorithm>

namespace mcrx {
  template <typename> class grid_cell;
  template <typename> class grid_base;
  template <typename, typename> class dynamic_grid;
  template <typename> class octogrid;
  template <typename> class adaptive_grid;
  template<typename> class grid_factory;
  template<typename> class grid_creator;
}

namespace CCfits {
  class FITS;
  class ExtHDU;
}


#include "grid_cell.h"

// *** abstract base class grid_base ***

/** Abstract base class for all grids, containing fundamental
     quantities. The grid_base contains parameters such as the extent
     of the grid, the size of the cells and a pointer to the parent
     cell in a hierarchichal grid.  It is derived by the concrete grid
     classes octogrid and dynamic_grid.

     The grid_base also defines iterators (to be implemented in the
     concrete classes) that allows for traversal of the grid cells,
     both in a depth-first manner through the entire hierarchical grid
     structure and locally in this grid only. The iterators are
     STL-compliant input iterators. */
template < typename cell_data_type >
class mcrx::grid_base {
public:
  typedef mcrx::T_float T_float;
  typedef grid_base T_grid;
  typedef grid_cell<cell_data_type> T_cell;
  typedef typename T_cell::cell_tracker T_cell_tracker;
  typedef cell_data_type T_data;
  friend class grid_cell<cell_data_type> ;
  template<typename cdt> friend class adaptive_grid;

private:
  // undefined to prevent implicit generation
  grid_base (const T_grid&);
  T_grid& operator= (const T_grid&);
  template <typename TT_grid,typename TT_cell> class iterator_impl;

protected:
  // Grid dimensions
  vec3d min; ///< Minimum point of the grid cube.
  vec3d max; ///< Maximum point of the grid cube.
  vec3d size; ///< Size of grid cells.
  vec3d invsize; ///< 1/size, precomputed for efficiency.

  T_cell* parent; ///< Pointer to parent cell.

  virtual T_grid* create_sub (T_cell *) const = 0;
  virtual T_grid* create_sub (T_cell *, void*&) const = 0;
  virtual const T_cell* index (const vec3d &,
			       bool accept_outside = false) const;
  const T_cell*locate_samelevel (const vec3d &, int) const ;
  const T_cell* parent_cell () const {return parent;};
  T_cell* parent_cell () {return parent;};
  // because I don't know how to make the derived classes friends of
  // grid_cell
  void assign_cell (T_cell & lhs, const T_cell& rhs) const {
    lhs = rhs;};
  void load_recursive (std::vector<bool>::const_iterator&,
		       const std::vector<bool>::const_iterator&,
		       void*&);

  void placement_new_destructor ();

public:
  // creates a grid_base with specified dimensions and no parent cell
  grid_base (const vec3d & mi, const vec3d & ma, const vec3i & nn):
    min (mi), max (ma),
    size ((ma-mi)/nn), invsize (1./size),
    parent (0) 
  {ASSERT_ALL(size>0);};
  // creates a grid_base from the specified parent cell
  grid_base (T_cell *p, const vec3i & nn):
    min (p->getmin()), max (p->getmax()),
    size ((max-min)/nn),
    invsize (1/size), parent (p)
  {ASSERT_ALL(size>0);};
  // creates a grid_base with information from a gridstructure HDU
  grid_base (CCfits::ExtHDU& file, const vec3d& =vec3d(0,0,0));
  virtual ~grid_base () {};

  // the cell retrieval function has to be virtual
  virtual T_cell *get_cell (const typename T_cell::T_index& indx) = 0;
  virtual const T_cell *get_cell (const typename T_cell::T_index & indx) const = 0;

  /// Returns min point of the grid cube.
  const vec3d& getmin() const {return min;};
  /// Returns max poitn of the grid cube.
  const vec3d& getmax() const {return max;};
  /// Returns the size of the grid cells.
  const vec3d& getsize() const {return size;};
  /// Returns the inverse size of the grid cells (precomputed for efficiency.)
  const vec3d& getinvsize() const {return invsize;};
  /// Returns the number of cells in this grid.
  virtual const vec3i& getn () const = 0;
  /** Returns the number of cells in the subgrids created when a
      grid_cell is refined. This function is a little subtle, the only
      way I know how to extract this information is to actually create
      a grid... */
  vec3i get_subdiv () const {
    T_cell c (const_cast<T_grid*> (this), typename T_cell::T_index (0, 0, 0));
    c.refine();return c.sub_grid()->getn();};

  /// Identifies the grid cell containing a point.
  virtual T_cell_tracker locate (const vec3d & p, 
				 bool accept_outside = false) const;

  /** Checks if the ray hits the grid.  Returns the distance between
      the ray position and the point where the ray enters the grid and
      the cell it will enter.  If the ray misses, NaN and an invalid
      tracker is returned. If the ray is already inside the grid, we
      return the cell it's in and NaN as the length. */
  std::pair<T_cell_tracker, T_float> intersection_from_without (const ray_base& ray) {
    const T_float l = ray.intersect_cube(getmin(), getmax(), false);
    if(l==l)
      // outside
      return std::make_pair(locate(ray.position() + l*ray.direction(), true), l); 
    else
      // inside
      return std::make_pair(locate(ray.position(), false), l);
  };

  void statistics_recursive (std::vector<int>&, int) const;
  void get_structure (std::vector<bool>& structure,
		      std::vector<vec3i>* subdiv=0) const;

  // The iterator classes are STL-compliant input iterators:
  
  // Iterator that goes down into sub grids to find leaf cells.
  class iterator;
  class const_iterator;
  // The local iterators don't care about sub grids but only traverse
  // THIS grid, and return non-leaf cells.
  class local_iterator;
  class const_local_iterator;
 
  virtual const_iterator begin () const = 0;
  virtual iterator begin () = 0;
  virtual const_local_iterator lbegin () const = 0;
  virtual local_iterator lbegin () = 0;
  virtual const_iterator end () const = 0;
  virtual iterator end () = 0;
  virtual const_local_iterator lend () const = 0;
  virtual local_iterator lend () = 0;
};

#include "grid_iterator.h"
#include "grid_impl.h"


// *** top-level class adaptive_grid ***

/** Class representing the top level of a hierarchical grid
    structure. This is a special class because there are some
    operations that only make sense from the top level. Most grid
    operations are just passed down to the actual grid object. */
template <typename cell_data_type>
class mcrx::adaptive_grid {
public:
  typedef mcrx::dynamic_grid< cell_data_type, 
			      mcrx::octogrid<cell_data_type> > grid_type;
  typedef mcrx::T_float T_float;
  typedef grid_base<cell_data_type> T_grid;
  typedef grid_cell<cell_data_type> T_cell;
  typedef typename T_cell::cell_tracker T_cell_tracker;
  typedef cell_data_type T_data;
  friend class grid_cell<cell_data_type> ;
  friend class grid_creator<adaptive_grid< cell_data_type> >;
private:
  // undefined to prevent implicit generation
  adaptive_grid (const T_grid&);
  adaptive_grid& operator= (const adaptive_grid&);
  void structure_plot_point (std::ostream&,const vec3d&) const;

  /// Helper function that creates subgrids according to the supplied
  /// structure vector.
  void create_structure(const std::vector<bool>& structure);

protected:
  grid_type g_; ///< A grid is the only data member in this class.

  /// Memory buffer for the allocation of sub grids, if used.
  void* subgrid_block;

  typedef typename grid_type::local_iterator local_iterator;
  typedef typename grid_type::const_local_iterator const_local_iterator;
  const_local_iterator lbegin () const {return g_.lbegin ();};
  local_iterator lbegin () {return g_.lbegin ();};
  const_local_iterator lend () const {return g_.lend ();};
  local_iterator lend () {return g_.lend ();};

public:
  /// Creates an adaptive_grid with specified dimensions and no parent cell.
  adaptive_grid (const vec3d & mi, const vec3d & ma, const vec3i & nn):
    g_(mi, ma, nn), subgrid_block (0) {};

  /// Creates an adaptive_grid with information from a structure vector.
  adaptive_grid (const vec3d & mi, const vec3d & ma, const vec3i & nn,
		 const std::vector<bool>& structure);
  /// Creates an adaptive_grid with information from a gridstructure HDU.
  adaptive_grid (CCfits::ExtHDU& file, 
		 const vec3d& =vec3d(0,0,0));
  // creates an adaptive grid with specified base dimensions and uses
  // the grid_factory object to create the grid structure and cell data.
  adaptive_grid (const vec3d & mi, const vec3d & ma, const vec3i & nn,
		 grid_factory<T_data>& f);
  virtual ~adaptive_grid () {
    if (subgrid_block) {
      // Go through the hierarchy and call destructors, not delete
      g_.placement_new_destructor();
      operator delete (subgrid_block);}};

  // simple data retrieval functions
  const vec3d& getmin() const {return g_.getmin ();};
  const vec3d& getmax() const {return g_.getmax() ;};
  const vec3i& get_base_n() const {return g_.getn();};
  vec3i get_subdiv () const {return g_.get_subdiv ();};

  // cell location
  T_cell_tracker locate (const vec3d &p, bool accept_outside = false ) {
    return g_.locate (p, accept_outside);};

  std::pair<T_cell_tracker, T_float> intersection_from_without (const ray_base& ray) {
    return g_.intersection_from_without (ray);};

  // outputs grid structure to a stream
  void structure_plot (std::ostream&) const;
  /** Returns a structure vector for the grid. */
  std::vector<bool> get_structure () const {
    std::vector<bool> structure;
    g_.get_structure(structure);
    return structure;
  };

  /// saves grid structure to a FITS HDU (include grid-fits.h)
  void save_structure (CCfits::FITS& file, const std::string& hdu,
		       const std::string& length_unit = "") const;
  /// Returns the number of cells in the grid.
  int n_cells () const;
  /// Returns the number in the grid enumeration of the specified cell.
  int get_cell_number(const T_cell_tracker& c) const;

  /** Returns statistics about the number of cells at different
      refinement levels in the grid. */
  std::vector<int> statistics () const;

  // We just expose the iterators from grid_base
  typedef typename grid_type::iterator iterator;
  typedef typename grid_type::const_iterator const_iterator;
  const_iterator begin () const {return g_.begin ();};
  iterator begin () {return g_.begin ();};
  const_iterator end () const {return g_.end ();};
  iterator end () {return g_.end ();};

  bool check_integrity() const;
};



// *** grid_base methods ***

/** Constructor creates grid from a structure vector. When
    execution enters this constructor, the grid_base constructor has
    completed. This function then looks for the refinement mechanism
    and ensures that the file is consistent with the class.  A memory
    block is allocated to hold all the subgrids, so that they are well
    localized in memory, and finally load_recursive is called to
    actually create the structure. */
template < typename cell_data_type>
mcrx::adaptive_grid<cell_data_type>::
  adaptive_grid (const vec3d & mi, const vec3d & ma, const vec3i & nn,
		 const std::vector<bool>& structure) :
    g_(mi, ma, nn), subgrid_block (0)
{
  // currently this grid does not make use of subdivisor, so it's in
  // principle possible to make it crash by feeding a structure vector
  // that is incompatible with the grid class. However, since we so
  // far have not implemented anything but the octogrid subgrid, that
  // seems like a low risk. For now we hardcode a check.
  vec3d subdiv(2,2,2);
  ASSERT_ALL (g_.get_subdiv () == subdiv);

  create_structure(structure);
}

template < typename cell_data_type>
void mcrx::adaptive_grid<cell_data_type>::
create_structure(const std::vector<bool>& structure)
{   
  // Figure out how many sub grids we have (we have one for every
  // true in the structure vector)
  int n_subgrids = std::count (structure.begin(), structure.end(), true);
  const size_t block_size = n_subgrids*sizeof(typename grid_type::T_subgrid); 
  std::cout << "Allocating a memory block for " << n_subgrids
	    << " subgrids, " << block_size*1.0/(1024*1024) << " MB."  << std::endl;
  subgrid_block = operator new (block_size);
  
  std::vector<bool>::const_iterator i = structure.begin();
  void* sb = subgrid_block;
  g_.load_recursive (i, structure.end(), sb);
  assert(i==structure.end()); // make sure we used exactly all elements
}


/** Kind-of destructor.  This interesting construct is necessary since
 if we have allocated the grid using placement new, we can't let the
 destructors call delete.  So, we have to go through the data before,
 explicitly call the destructor for the substructure, set the pointer
 to 0 and THEN destruct ourselves.

 Note that this should really be a virtual function defined in the
 implementation grids, because if any of them were to have pointers,
 they would have to explicitly call destructors for them as well. */
template <typename cell_data_type>
void mcrx::grid_base<cell_data_type>::placement_new_destructor ()
{
  for (local_iterator c = lbegin (); c != lend (); ++c) {
    c->placement_new_destructor();
  }
}


/** Calculate the index of the cell at position p.  If p is outside of
 the grid, and there is a higher-level grid, calculate the index of
 the closest cell. Note that this is a virtual function, and this is a
 default implementation.  Octogrid redefines this because it can be
 much more efficient. */
template < typename cell_data_type >
const typename mcrx::grid_base<cell_data_type>::T_cell*
mcrx::grid_base<cell_data_type>::index (const vec3d& p,
					bool accept_outside) const
{
  //  calculate "floating point index" (it can be between -1 and 0 so
  //  we have to use floating points to detect this since it would
  //  round to 0.
  const vec3d iii=(p- getmin() )*getinvsize ();

  // Check for out of bounds
  bool err= false;

  typename T_cell::T_index ind;

  if(iii[0]<0) {
    ind [0]=0;
    err= true;
  }
  else if(iii[0]>=getn ()[0]) {
    ind[0]=getn ()[0]-1;
    err= true;
  }
  else
    ind [0] = static_cast<typename T_cell::T_indextype> (iii [0]);

  if(iii[1]<0) {
    ind [1]=0;
    err= true;
  }
  else if(iii[1]>=getn ()[1]) {
    ind[1]=getn ()[1]-1;
    err= true;
  }
  else
    ind [1] = static_cast<typename T_cell::T_indextype> (iii [1]);

  if(iii[2]<0) {
    ind [2]=0;
    err= true;
  }
  else if(iii[2]>=getn ()[2]) {
    ind[2]=getn ()[2]-1;
    err= true;
  }
  else
    ind [2] = static_cast<typename T_cell::T_indextype> (iii [2]);
  
  // If we are in a top-level grid, return null for out-of-bonds
  // (unless we have explicitly asked for out with accept_outside) 
  if(err && !accept_outside && !parent_cell ())
    return 0;
  else
    return get_cell (ind);
}


/** Find the grid_cell that contains position p.  If this grid is not
 in a top-level grid, we accept positions outside of the grid itself,
 and just use the closest coordinates. This is used by neighbor. */
template < typename cell_data_type >
typename mcrx::grid_base<cell_data_type>::T_cell_tracker
mcrx::grid_base<cell_data_type>::locate(const vec3d& p,
					bool accept_outside) const
{
  const T_cell* const c(index (p, accept_outside));
  
  if(c && !c->is_leaf())  
    // That cell exists but is subdivided, recurse in the locate call
    // (accept_outside is not necessary in this case because it only
    // affects the top-level grid)
    return c->sub_grid()->locate(p);
  else
    // cell is not subdivided, just return it
    return const_cast<T_cell*>(c);
}

/** Finds the grid_cell containing position p at a specific
    hierarchical level in the grid. This routine works like the normal
    locate function, except that it goes down a maximum of level
    levels. It is used by the neighbor_samelevel function in the
    grid_cell class. */
template < typename cell_data_type >
const typename mcrx::grid_base<cell_data_type>::T_cell*
mcrx::grid_base<cell_data_type>::locate_samelevel (const vec3d& p,
						   int level) const
{
  const T_cell* const c = index (p);
  // should never get zero here
  assert(c);
  
  if(c && (!c->is_leaf()) && (level > 0))
    // That cell exists but is subdivided, recurse in the locate call
    return c->sub_grid()->locate_samelevel (p, level - 1);
  else
    // cell is not subdivided, just return it
    return c;
}

template <typename cell_data_type>
int mcrx::adaptive_grid<cell_data_type>::n_cells () const
{
  int n = 0;
  for (const_iterator c = begin (), e = end (); c != e; ++c, ++n);
  return n;
}

template <typename cell_data_type>
std::vector<int>
mcrx::adaptive_grid<cell_data_type>::statistics () const
{
  std::vector<int> stats;
  g_.statistics_recursive (stats, 0);
  return stats;
}

template <typename cell_data_type>
int mcrx::adaptive_grid<cell_data_type>::
get_cell_number (const T_cell_tracker& c) const
{
  int n = 0;
  for (const_iterator i = begin (), e = end (); i != e; ++i, ++n)
    if (&(*i)==c.cell())
      return n;
  throw 1;
}

/** Recursively collects refinement statistics. */
template < typename cell_data_type >
inline void
mcrx::grid_base<cell_data_type>::statistics_recursive (std::vector<int>& stats,
						       int level) const
{
  // increase size if we have to
  if (stats.size() <level + 1)
    stats.push_back(0 );
  for (const_local_iterator c = lbegin (), e = lend (); c != e; ++c) {
    if (c->is_leaf())
      ++ stats [level];
    else
      c->sub_grid()->statistics_recursive(stats, level + 1 ); 
  }
}

/** Collects grid structure information when saving structure to a
    file.  This function recursively calls itself on all cells in this
    grid, saving subdivision information in the vectors. */
template<typename cell_data_type>
void
mcrx::grid_base<cell_data_type>::get_structure (std::vector<bool>& structure,
						std::vector<vec3i>* subdivision) const
{
  for (const_local_iterator i = lbegin (); i != lend (); ++i) {
    structure.push_back(!i->is_leaf());
    if (subdivision && !i->is_leaf())
      subdivision->push_back(i->sub_grid()->getn() );
    if (!i->is_leaf())
      i->sub_grid()->get_structure(structure, subdivision );
  }
}


template <typename cell_data_type>
void
mcrx::adaptive_grid<cell_data_type>::
structure_plot_point (std::ostream& o, const vec3d& p) const
{
    o << p [0]  <<  '\t'
     << p [1]  <<  '\t'
      << p [2]  <<  '\n';
}

/** Saves the grid structure as a text file to the specified
    stream. The structure consists of a bunch of line segments that
    you can feed to e.g. gnuplot to plot the grid. */
template <typename cell_data_type>
void 
mcrx::adaptive_grid<cell_data_type>::
structure_plot (std::ostream& o) const
{
  o.precision (12); 
  typename T_grid::const_iterator j = begin();
  while (j != end()) {
    const vec3d a = j->getmin(); 
    const vec3d b = j->getmax ();
    const vec3d p1 = a;
    const vec3d p2 (b [0], a [1], a [2]);
    const vec3d p3 (b [0], b [1], a [2]);
    const vec3d p4 (a [0], b [1], a [2]);
    const vec3d p5 (a [0], a [1], b [2]);
    const vec3d p6 (b [0], a [1], b [2]);
    const vec3d p7 (b [0], b [1], b [2]);
    const vec3d p8 (a [0], b [1], b [2]);

    structure_plot_point (o, p1);
    structure_plot_point (o, p2);
    structure_plot_point (o, p3);
    structure_plot_point (o, p4);
    structure_plot_point (o, p1);
    structure_plot_point (o, p5);
    structure_plot_point (o, p6);
    structure_plot_point (o, p7);
    structure_plot_point (o, p8);
    structure_plot_point (o, p5);
    o << "\n";
    structure_plot_point (o, p2);
    structure_plot_point (o, p6);
    o << "\n";
    structure_plot_point (o, p4);
    structure_plot_point (o, p8);
    o << "\n";
    structure_plot_point (o, p3);
    structure_plot_point (o, p7);
    o << "\n";
   ++j;
  }
}

/// Checks that all leaf cells have a data object attached.
template <typename cell_data_type>
bool mcrx::adaptive_grid<cell_data_type>::check_integrity () const
{
  bool bad=false;
  for (const_iterator c = begin (), e = end (); c!=e; ++c) {
    bad = bad | (!c->data());
  }

  if (bad)
    std::cerr << "Grid integrity check failed!" << std::endl;

  return bad;    
}

#endif
